package priv.pfz.basedao.apt;

import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import priv.pfz.basedao.annotations.Column;
import priv.pfz.basedao.annotations.Entity;
import priv.pfz.basedao.annotations.NoColumn;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Filer;
import javax.annotation.processing.Messager;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeKind;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import javax.tools.Diagnostic;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

/**
 * @author yyjxp
 * 2024/1/28 19:14
 */
public class AnnotationProcessor extends AbstractProcessor {

    private static final boolean DEBUG_LOG = true;

    private static final DateTimeFormatter DTF = DateTimeFormatter.ofPattern("yyyy/M/d H:m");

    private static final Map<String, String> COLUMN_TYPE_MAP = new HashMap<String, String>() {
        {
            put("java.lang.String", "String");
            put("java.lang.Integer", "Int");
            put("int", "Int");
            put("java.lang.Long", "Long");
            put("long", "Long");
            put("java.util.Date", "Date");
        }
    };

    private Types typeUtils;
    private Elements elementUtils;
    private Messager messager;
    private Filer filer;
    private String queryTemplate;
    private String columnTemplate;
    private String baseColumnTemplate;

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        Set<String> set = new HashSet<>();
        set.add(Entity.class.getName());
        return set;
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.RELEASE_8;
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        messager = processingEnv.getMessager();
        info("basedao开始自动生成Query类源码");

        //初始化
        typeUtils = processingEnv.getTypeUtils();
        elementUtils = processingEnv.getElementUtils();
        filer = processingEnv.getFiler();

        //读取模板
        queryTemplate = readResource("query_template");
        columnTemplate = readResource("column_template");
        baseColumnTemplate = readResource("base_column_template");

        //轮空处理
        if (annotations.isEmpty()) {
            return true;
        }
        TypeElement entityAnnotation = annotations.iterator().next();

        //根据Entity生成Query
        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(entityAnnotation);
        for (Element element : elements) {
            if (element.getKind() != ElementKind.CLASS) {
                info(element.getSimpleName() + "类型非法,跳过");
                continue;
            }
            TypeElement typeElement = (TypeElement) element;
            info("开始处理" + typeElement.getQualifiedName());
            SourceFile sourceFile = generateQuery(typeElement);
            writeSourceFile(sourceFile.getFilename(), typeElement, sourceFile.getContent());
        }

        return false;
    }

    private void info(String str) {
        messager.printMessage(Diagnostic.Kind.NOTE, str);
    }

    private void debug(String str) {
        if (DEBUG_LOG) {
            messager.printMessage(Diagnostic.Kind.NOTE, str);
        }
    }

    private String readResource(String filename) {
        try (InputStream is = filer.getResource(StandardLocation.CLASS_PATH, "", filename).openInputStream()) {
            return IOUtils.toString(is, StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new RuntimeException("读取资源文件异常:" + filename, e);
        }
    }

    private void writeSourceFile(String filename, Element oriElement, String content) {
        try (OutputStream os = filer.createSourceFile(filename, oriElement).openOutputStream()) {
            IOUtils.write(content, os, StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new RuntimeException("写java源文件异常:" + filename, e);
        }
    }

    /**
     * 根据Entity类生成Query类
     */
    private SourceFile generateQuery(TypeElement entityClass) {
        String entitySuffix = entityClass.getAnnotation(Entity.class).suffix();
        String qualifiedName = entityClass.getQualifiedName().toString();
        String pkg = qualifiedName.substring(0, qualifiedName.lastIndexOf('.'));
        String queryName = entityClass.getSimpleName().toString().replace(entitySuffix, "") + "Query";

        //构造query中的column
        List<String> columns = new ArrayList<>();
        Set<String> imports = new LinkedHashSet<>();
        List<AptFieldInfo> fieldInfos = getFieldInfos(entityClass);
        for (AptFieldInfo fieldInfo : fieldInfos) {
            String fieldTypeName = fieldInfo.getFieldType().toString();
            String column;
            if (COLUMN_TYPE_MAP.containsKey(fieldTypeName)) {
                //常见类型
                Map<String, String> params = new HashMap<>();
                params.put("type", COLUMN_TYPE_MAP.get(fieldTypeName));
                params.put("classname", queryName);
                params.put("fieldName", fieldInfo.getFieldName());
                params.put("columnName", fieldInfo.getColumnName());
                column = AptUtils.fillTemplate(columnTemplate, params);
            } else {
                //其它类型
                Map<String, String> params = new HashMap<>();
                params.put("type", fieldTypeName.substring(fieldTypeName.lastIndexOf('.') + 1));
                params.put("classname", queryName);
                params.put("fieldName", fieldInfo.getFieldName());
                params.put("columnName", fieldInfo.getColumnName());
                column = AptUtils.fillTemplate(baseColumnTemplate, params);
                imports.add("import " + fieldTypeName + ";");
            }
            columns.add(column);
        }

        //构造query
        Map<String, String> params = new HashMap<>();
        params.put("package", pkg);
        params.put("imports", String.join("\n", imports));
        params.put("date", LocalDateTime.now().format(DTF));
        params.put("classname", queryName);
        params.put("columns", String.join("\n", columns));
        String query = AptUtils.fillTemplate(queryTemplate, params);

        String filename = pkg + "." + queryName;
        return new SourceFile(filename, query);
    }

    /**
     * 解析Entity类的属性
     * 1、serialVersionUID字段不计入
     * 2、static修饰的字段不计入
     * 3、transient修饰的字段不计入
     * 4、@NoColumn标注的字段不计入
     * 5、支持父类中的字段（子类覆盖父类的情况，以子类为准）
     */
    private List<AptFieldInfo> getFieldInfos(TypeElement entityClass) {
        Set<String> handledSet = new HashSet<>();
        List<AptFieldInfo> fieldInfos = new ArrayList<>();
        while (true) {
            for (Element element : entityClass.getEnclosedElements()) {
                if (element.getKind() != ElementKind.FIELD) {
                    continue;
                }
                VariableElement varElement = (VariableElement) element;
                String fieldName = varElement.getSimpleName().toString();
                if (!handledSet.add(fieldName)) {
                    continue;
                }
                if (fieldName.equals("serialVersionUID")) {
                    continue;
                }
                if (varElement.getModifiers().contains(Modifier.STATIC)) {
                    continue;
                }
                if (varElement.getModifiers().contains(Modifier.TRANSIENT)) {
                    continue;
                }
                if (varElement.getAnnotation(NoColumn.class) != null) {
                    continue;
                }
                AptFieldInfo fieldInfo = new AptFieldInfo();
                fieldInfo.setFieldName(fieldName);
                fieldInfo.setColumnName(getColumnName(varElement, fieldName));
                fieldInfo.setFieldType(varElement.asType());
                fieldInfos.add(fieldInfo);
            }
            if (entityClass.getSuperclass().getKind() == TypeKind.NONE) {
                break;
            }
            entityClass = (TypeElement) typeUtils.asElement(entityClass.getSuperclass());
        }
        return fieldInfos;
    }

    private String getColumnName(VariableElement varElement, String fieldName) {
        Column column = varElement.getAnnotation(Column.class);
        if (column != null && StringUtils.isNotBlank(column.columnName())) {
            return column.columnName();
        }
        return AptUtils.fieldToColumn(fieldName);
    }
}
